# coding: utf-8

import pandas as pd
import os
import pytz
from datetime import datetime
import pymongo
import json
from time import time
import MySQLdb
import types

"""
把1分钟resample成5分钟，15分钟，30分钟和60分钟,120分钟
start:起始日期
end:截至日期
"""


def _get_stocks_from_db():
    # connect to mysql
    conn = MySQLdb.connect(host='rdsqtehrv8tqh7v60yvujpublic.mysql.rds.aliyuncs.com',
                           user='tradingfloor',
                           passwd='tradingfloor123',
                           db='tradingreason',
                           charset='utf8')
    # conn = MySQLdb.connect(host='localhost',
    #                   user='root',
    #                   passwd='root',
    #                   db='tradingreasontt',
    #                   charset='utf8')
    cur = conn.cursor()
    res = []
    cur.execute("select symbol from stock_info")
    for (symbol,) in cur:
        res.append(symbol)
    return res


t0 = time()

tz = pytz.timezone(pytz.country_timezones('cn')[0])
start = datetime(2017, 6, 29, 9, 0, tzinfo=tz)
end = datetime(2017, 6, 30, 16, 0, tzinfo=tz)

log = open('resample_log.txt', 'a')
td = datetime.today()
log.write('#####################################################\n')
log.write('task begins at: ' + td.strftime('%Y-%m-%d %H:%M:%S') + '\n')
log.write('data range:\n')
log.write('start date :' + start.strftime('%Y-%m-%d %H:%M:%S') + '\n')
log.write('end date :' + end.strftime('%Y-%m-%d %H:%M:%S') + '\n')

client = pymongo.MongoClient('121.40.212.219', 27017)  # '121.40.212.219', 27017
db = client.Minute_data

seclist = _get_stocks_from_db()  # ['SZ300628']

for sec in seclist:
    cur = db.one_min_data.find({'symbol': sec, 'dateTime': {'$gte': start, '$lte': end}}, {'_id': 0}) \
        .sort("dateTime", pymongo.ASCENDING)
    df = pd.DataFrame(list(cur))
    if df.empty:
        if type(sec) != types.NoneType:
            print(sec + 'is empty')
            log.write(sec + ' is empty' + '\n')
        else:
            print("sec is none")
        continue
    log.write('generating: ' + sec + '\n')

    print(sec)
    df = df.set_index('dateTime')
    # ---------------write in 5 minute data----------------
    df_m5 = pd.DataFrame([])
    df_m5['open'] = df['open'].resample('5min', how='first', closed='right', label='right').dropna()
    df_m5['close'] = df['close'].resample('5min', how='last', closed='right', label='right').dropna()
    df_m5['high'] = df['high'].resample('5min', how='max', closed='right', label='right').dropna()
    df_m5['low'] = df['low'].resample('5min', how='min', closed='right', label='right').dropna()
    df_m5['barTime'] = df['barTime'].resample('5min', how='last', closed='right', label='right').dropna()
    df_m5['date'] = df['date'].resample('5min', how='last', closed='right', label='right').dropna()
    df_m5['amount'] = df['amount'].resample('5min', how='sum', closed='right', label='right').dropna()
    df_m5['volume'] = df['volume'].resample('5min', how='sum', closed='right', label='right').dropna()
    df_m5['symbol'] = sec
    db.min_5_data.insert_many(df_m5.reset_index().to_dict(orient='records'))
    print("5 min done")

    # ---------------write in 15 minute data----------------
    df_m15 = pd.DataFrame([])
    df_m15['open'] = df_m5['open'].resample('15min', how='first', closed='right', label='right').dropna()
    df_m15['close'] = df_m5['close'].resample('15min', how='last', closed='right', label='right').dropna()
    df_m15['high'] = df_m5['high'].resample('15min', how='max', closed='right', label='right').dropna()
    df_m15['low'] = df_m5['low'].resample('15min', how='min', closed='right', label='right').dropna()
    df_m15['barTime'] = df_m5['barTime'].resample('15min', how='last', closed='right', label='right').dropna()
    df_m15['date'] = df_m5['date'].resample('15min', how='last', closed='right', label='right').dropna()
    df_m15['amount'] = df_m5['amount'].resample('15min', how='sum', closed='right', label='right').dropna()
    df_m15['volume'] = df_m5['volume'].resample('15min', how='sum', closed='right', label='right').dropna()
    df_m15['symbol'] = sec
    db.min_15_data.insert_many(df_m15.reset_index().to_dict(orient='records'))
    print("15 min done")

    # ---------------write in 30 minute data----------------
    df_m30 = pd.DataFrame([])
    df_m30['open'] = df_m15['open'].resample('30min', how='first', closed='right', label='right').dropna()
    df_m30['close'] = df_m15['close'].resample('30min', how='last', closed='right', label='right').dropna()
    df_m30['high'] = df_m15['high'].resample('30min', how='max', closed='right', label='right').dropna()
    df_m30['low'] = df_m15['low'].resample('30min', how='min', closed='right', label='right').dropna()
    df_m30['barTime'] = df_m15['barTime'].resample('30min', how='last', closed='right', label='right').dropna()
    df_m30['date'] = df_m15['date'].resample('30min', how='last', closed='right', label='right').dropna()
    df_m30['amount'] = df_m15['amount'].resample('30min', how='sum', closed='right', label='right').dropna()
    df_m30['volume'] = df_m15['volume'].resample('30min', how='sum', closed='right', label='right').dropna()
    df_m30['symbol'] = sec
    db.min_30_data.insert_many(df_m30.reset_index().to_dict(orient='records'))
    print("30 min done")

    # ---------------write in 60 minute data----------------
    # utc time 1:30 to 3:30
    df_m60_morning = pd.DataFrame([])
    df_m60_morning['open'] = df_m30['open'].between_time('1:30', '3:30'). \
        resample('60min', how='first', closed='right', label='right', base=30).dropna()

    df_m60_morning['close'] = df_m30['close'].between_time('1:30', '3:30'). \
        resample('60min', how='last', closed='right', label='right', base=30).dropna()

    df_m60_morning['high'] = df_m30['high'].between_time('1:30', '3:30'). \
        resample('60min', how='max', closed='right', label='right', base=30).dropna()

    df_m60_morning['low'] = df_m30['low'].between_time('1:30', '3:30'). \
        resample('60min', how='min', closed='right', label='right', base=30).dropna()

    df_m60_morning['barTime'] = df_m30['barTime'].between_time('1:30', '3:30'). \
        resample('60min', how='last', closed='right', label='right', base=30).dropna()

    df_m60_morning['date'] = df_m30['date'].between_time('1:30', '3:30'). \
        resample('60min', how='last', closed='right', label='right', base=30).dropna()

    df_m60_morning['amount'] = df_m30['amount'].between_time('1:30', '3:30'). \
        resample('60min', how='sum', closed='right', label='right', base=30).dropna()

    df_m60_morning['volume'] = df_m30['volume'].between_time('1:30', '3:30'). \
        resample('60min', how='sum', closed='right', label='right', base=30).dropna()
    df_m60_morning['symbol'] = sec
    db.min_60_data.insert_many(df_m60_morning.reset_index().to_dict(orient='records'))
    # utc time 5:00 to 7:00
    df_m60_afternoon = pd.DataFrame([])
    df_m60_afternoon['open'] = df_m30['open'].between_time('5:00', '7:00'). \
        resample('60min', how='first', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['close'] = df_m30['close'].between_time('5:00', '7:00'). \
        resample('60min', how='last', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['high'] = df_m30['high'].between_time('5:00', '7:00'). \
        resample('60min', how='max', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['low'] = df_m30['low'].between_time('5:00', '7:00'). \
        resample('60min', how='min', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['barTime'] = df_m30['barTime'].between_time('5:00', '7:00'). \
        resample('60min', how='last', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['date'] = df_m30['date'].between_time('5:00', '7:00'). \
        resample('60min', how='last', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['amount'] = df_m30['amount'].between_time('5:00', '7:00'). \
        resample('60min', how='sum', closed='right', label='right', base=0).dropna()

    df_m60_afternoon['volume'] = df_m30['volume'].between_time('5:00', '7:00'). \
        resample('60min', how='sum', closed='right', label='right', base=0).dropna()
    df_m60_afternoon['symbol'] = sec
    db.min_60_data.insert_many(df_m60_afternoon.reset_index().to_dict(orient='records'))
    print("60 min done")

    # ---------------write in 120 minute data----------------
    # utc time 1:30 to 3:30
    df_m120_morning = pd.DataFrame([])
    df_m120_morning['open'] = df_m30['open'].between_time('1:30', '3:30'). \
        resample('120min', how='first', closed='right', label='right', base=90).dropna()
    df_m120_morning['close'] = df_m30['close'].between_time('1:30', '3:30'). \
        resample('120min', how='last', closed='right', label='right', base=90).dropna()
    df_m120_morning['high'] = df_m30['high'].between_time('1:30', '3:30'). \
        resample('120min', how='max', closed='right', label='right', base=90).dropna()
    df_m120_morning['low'] = df_m30['low'].between_time('1:30', '3:30'). \
        resample('120min', how='min', closed='right', label='right', base=90).dropna()
    df_m120_morning['barTime'] = df_m30['barTime'].between_time('1:30', '3:30'). \
        resample('120min', how='last', closed='right', label='right', base=90).dropna()
    df_m120_morning['date'] = df_m30['date'].between_time('1:30', '3:30'). \
        resample('120min', how='last', closed='right', label='right', base=90).dropna()
    df_m120_morning['amount'] = df_m30['amount'].between_time('1:30', '3:30'). \
        resample('120min', how='sum', closed='right', label='right', base=90).dropna()
    df_m120_morning['volume'] = df_m30['volume'].between_time('1:30', '3:30'). \
        resample('120min', how='sum', closed='right', label='right', base=90).dropna()
    df_m120_morning['symbol'] = sec
    db.min_120_data.insert_many(df_m120_morning.reset_index().to_dict(orient='records'))
    # utc time 5:00 to 7:00
    df_m120_afternoon = pd.DataFrame([])
    df_m120_afternoon['open'] = df_m30['open'].between_time('5:00', '7:00'). \
        resample('120min', how='first', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['close'] = df_m30['close'].between_time('5:00', '7:00'). \
        resample('120min', how='last', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['high'] = df_m30['high'].between_time('5:00', '7:00'). \
        resample('120min', how='max', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['low'] = df_m30['low'].between_time('5:00', '7:00'). \
        resample('120min', how='min', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['barTime'] = df_m30['barTime'].between_time('5:00', '7:00'). \
        resample('120min', how='last', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['date'] = df_m30['date'].between_time('5:00', '7:00'). \
        resample('120min', how='last', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['amount'] = df_m30['amount'].between_time('5:00', '7:00'). \
        resample('120min', how='sum', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['volume'] = df_m30['volume'].between_time('5:00', '7:00'). \
        resample('120min', how='sum', closed='right', label='right', base=60).dropna()
    df_m120_afternoon['symbol'] = sec
    db.min_120_data.insert_many(df_m120_afternoon.reset_index().to_dict(orient='records'))
    print("120 min done")

    log.write('successfully write in: ' + sec + '\n')

log.write('task finishes in %.2fs\n' % (time() - t0))
log.close()
print('task finishes in %.2fs\n' % (time() - t0))
